#!/usr/bin/python

from datetime import date
from datetime import datetime
from datetime import timedelta


class ScheduleBlock(object):

  def __init__(self, start_time, duration):
    self.start_time = start_time
    self.duration = duration
    # TODO Number of doctors required?


class Schedule(object):

  def __init__(self, schedule_blocks):
    self.schedule_blocks = schedule_blocks
    self.days = len(self.schedule_blocks) // 2
    self.shifts = len(self.schedule_blocks)

  def ShiftCount(self):
    return self.shifts

  def AllShiftKeysArray(self):
    return range(self.shifts)

  def DayShiftKeysArray(self):
    return [i * 2 for i in range(self.shifts // 2)]

  def NightShiftKeysArray(self):
    return [i * 2 + 1 for i in range(self.shifts // 2)]

  def ScheduleBlocks(self):
    return self.schedule_blocks

  # TODO Distinguish between day and night shifts.
  # Swing shifts will make this more complicated.

  def IsDay(self, shift):
    return shift & 1 == 0

  def IsNight(self, shift):
    return shift & 1 == 1

  def ScheduleBlockForShift(self, shift):
    return self.schedule_blocks[shift]

  def ShiftNumbersForDate(self, day):
    return [(day - 1) * 2 - 1, (day - 1) * 2, (day - 1) * 2 + 1]

  def DayForShiftNumber(self, shift):
    return shift // 2 + 1


# Utility functions


def FirstDayOfMonthWithDelta(dt, delta_years=0, delta_months=0):
  # d_years, d_months are "deltas" to apply to dt
  y, m = dt.year + delta_years, dt.month + delta_months
  a, m = divmod(m - 1, 12)
  return datetime(y + a, m + 1, 1)


def CreateScheduleStartingAtMonth(start_time):
  schedule_blocks = []
  shift_start = start_time
  shift_end = FirstDayOfMonthWithDelta(start_time, delta_months=1)
  duration = timedelta(hours=12)
  while shift_start < shift_end:
    schedule_blocks.append(ScheduleBlock(shift_start, duration))
    shift_start = shift_start + duration
  return Schedule(schedule_blocks)
